#include <iostream>
#include <vector>
#include <random>
#include <string>
#include <cmath>
#include <cassert>
#include <tuple>


using namespace std;

// 全局参数
const int q = 1048576;  // 模数
const int n = 512;      // 矩阵维度
const int bar_n = 256;  // 错误矩阵列数
const int bar_m = 256;  // 错误矩阵行数

// 工具函数
int randomInt(int mod) {
    random_device rd;
    mt19937 gen(rd());
    uniform_int_distribution<> dis(0, mod - 1);
    return dis(gen);
}

vector<vector<int>> randomMatrix(int rows, int cols, int mod) {
    vector<vector<int>> matrix(rows, vector<int>(cols));
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            matrix[i][j] = randomInt(mod);
        }
    }
    return matrix;
}

vector<vector<int>> matrixAdd(const vector<vector<int>> &A, const vector<vector<int>> &B, int mod) {
    int rows = A.size();
    int cols = A[0].size();
    vector<vector<int>> C(rows, vector<int>(cols));
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            C[i][j] = (A[i][j] + B[i][j]) % mod;
        }
    }
    return C;
}

vector<vector<int>> matrixMul(const vector<vector<int>> &A, const vector<vector<int>> &B, int mod) {
    int rows = A.size();
    int cols = B[0].size();
    int k = A[0].size();
    vector<vector<int>> C(rows, vector<int>(cols, 0));
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            for (int l = 0; l < k; ++l) {
                C[i][j] = (C[i][j] + A[i][l] * B[l][j]) % mod;
            }
        }
    }
    return C;
}

// 哈希函数模拟
string hashFunc(const string &input) {
    return "hashed_" + input; // 简化实现，真实情况下需使用加密哈希
}

// KeyGen 算法
pair<pair<string, vector<vector<int>>>, tuple<string, string, vector<vector<int>>, vector<vector<int>>>> KeyGen() {
    // 生成随机种子
    string s = to_string(randomInt(q));
    string seed_E = to_string(randomInt(q));
    string seed_A = to_string(randomInt(q));

    // 使用随机种子 seed_A 生成随机矩阵 A
    auto A = randomMatrix(n, n, q);

    // 使用随机种子 seed_E 生成错误矩阵 S 和 E
    auto S = randomMatrix(n, bar_n, q);
    auto E = randomMatrix(n, bar_n, q);

    // 计算 B = AS + E
    auto B = matrixAdd(matrixMul(A, S, q), E, q);

    // 构造公钥 pk
    pair<string, vector<vector<int>>> pk = {seed_A, B};

    // 分步构造私钥 sk
    string sk_s = s;
    string sk_seed_A = seed_A;
    vector<vector<int>> sk_B = B;
    vector<vector<int>> sk_S = S;
    auto sk = make_tuple(sk_s, sk_seed_A, sk_B, sk_S);

    return {pk, sk};
}


// Encaps 算法
tuple<vector<vector<int>>, vector<vector<int>>, string, string> Encaps(const pair<string, vector<vector<int>>> &pk) {
    auto [seed_A, B] = pk;

    // 选择随机消息 m
    string m = to_string(randomInt(q));

    // 生成随机种子 r, k, d
    string r = to_string(randomInt(q));
    string k = to_string(randomInt(q));
    string d = to_string(randomInt(q));

    // 使用随机种子 r 生成错误矩阵 S' 和 E
    auto S_prime = randomMatrix(bar_m, n, q);
    auto E = randomMatrix(bar_m, n, q);
    auto E_double_prime = randomMatrix(bar_m, bar_n, q);

    // 生成矩阵 A
    auto A = randomMatrix(n, n, q);

    // 计算 B' = S'A + E
    auto B_prime = matrixAdd(matrixMul(S_prime, A, q), E, q);

    // 计算 V = S'B + E''
    auto V = matrixAdd(matrixMul(S_prime, B, q), E_double_prime, q);

    // 计算 C = V + Encode(m)（简化为直接加上 m 的数值表示）
    auto C = V;  // 真实实现需要对 m 进行编码

    // 计算 ss
    string ss = hashFunc(to_string(B_prime[0][0]) + to_string(C[0][0]) + k + d);

    return {B_prime, C, d, ss};
}

// Decaps 算法
string Decaps(const tuple<vector<vector<int>>, vector<vector<int>>, string, string> &cipher,
              const tuple<string, string, vector<vector<int>>, vector<vector<int>>> &sk) {
    // 解构 cipher
    auto B_prime = std::get<0>(cipher);
    auto C = std::get<1>(cipher);
    auto d = std::get<2>(cipher);
    auto ss = std::get<3>(cipher); // 如果 cipher 包含 4 个元素

    // 解构 sk
    auto s = std::get<0>(sk);
    auto seed_A = std::get<1>(sk);
    auto B = std::get<2>(sk);
    auto S = std::get<3>(sk);


    // 计算 D = C - B'S
    auto D = matrixAdd(C, matrixMul(B_prime, S, q), q);

    // Decode(D)（简化为直接还原为消息 m'）
    string m_prime = to_string(D[0][0]);

    // 生成随机种子 r', k', d'
    string r_prime = to_string(randomInt(q));
    string k_prime = to_string(randomInt(q));
    string d_prime = to_string(randomInt(q));

    // 验证条件
    if (d == d_prime) {
        return hashFunc(to_string(B_prime[0][0]) + to_string(C[0][0]) + k_prime + d);
    } else {
        return hashFunc(to_string(B_prime[0][0]) + to_string(C[0][0]) + s + d);
    }
}

int main() {
    // 生成公私钥
    auto [pk, sk] = KeyGen();

    // 执行密钥封装
    auto [B_prime, C, d, ss] = Encaps(pk);

    // 执行密钥解封
    string ss_dec = Decaps(std::make_tuple(B_prime, C, d, ss), sk);

    // 验证结果
    cout << "Shared secret key (Encaps): " << ss << endl;
    cout << "Shared secret key (Decaps): " << ss_dec << endl;

    if(ss == ss_dec) // 验证两者一致
    {
        cout << "验证通过" << endl;
    }
    return 0;
}
